﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Timtopia.PowerGraphBeamSearch
{
    using System.Collections;
    using System.Diagnostics;
    using System.IO;
    using Timtopia.DGMLPosterView;
    using Neighbours = Timtopia.DGMLPosterView.PowerGraphDecomposition.Neighbours<int>;
    using Tree = Timtopia.DGMLPosterView.PowerGraphDecomposition.TreeNode<int?>;

    public class Module : IEnumerable<int>
    {
        public Tree TreeNode;
        public HashSet<Module> Outgoing = new HashSet<Module>();
        public HashSet<Module> Incoming = new HashSet<Module>();
        public bool IsLeaf {
            get { return TreeNode.Leaf != null; }
        }

        public IEnumerator<int> GetEnumerator()
        {
            return TreeNode.Leaves.Select(l => l.Value).GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }
    }

    public class Modules
    {
        Dictionary<Tree, Module> treeModuleMap = new Dictionary<Tree, Module>();

        public static int EdgeWeight = 1;
        public static int CrossingWeight = 1;
        public static int ModuleWeight = 1;

        public int N { get { return treeModuleMap.Count; } }
        Tree tree;

        public Tree Tree
        {
            get { return tree; }
        }

        public Module GetModule(Tree t)
        {
            Module m = null;
            treeModuleMap.TryGetValue(t, out m);
            return m;
        }

        int cost = -1;
        int edges = -1;
        int crossings = -1;
        int leafCount;

        public int ModuleCount
        {
            get
            {
                return N - leafCount; // tree.Leaves.Count();
            }
        }

        public int EdgeCount
        {
            get
            {
                if (edges < 0)
                {
                    int c = Cost;
                }
                return edges;
            }
        }

        public int CrossingCount
        {
            get
            {
                if (crossings < 0)
                {
                    int c = Cost;
                }
                return crossings;
            }
        }

        public int Cost
        {
            get
            {
                if (cost < 0)
                {
                    edges = 0;
                    crossings = 0;
                    foreach (var u in treeModuleMap.Values)
                    {
                        foreach (var v in u.Outgoing)
                        {
                            if (CrossingWeight != 0)
                            {
                                var tu = u.TreeNode;
                                var tv = v.TreeNode;
                                Tree a, b;
                                var l = tu.LeastCommonAncestor(tv, out a, out b);
                                crossings += tu.PathToRoot(a).Count() + tv.PathToRoot(b).Count();
                            }
                            ++edges;
                        }
                    }
                    return cost = edges * EdgeWeight + crossings * CrossingWeight + ModuleCount * ModuleWeight;
                }
                return cost;
            }
        }

        private static Tree CopyTree(Tree orig, Dictionary<Tree, Tree> nodeMap)
        {
            Tree copy = new Tree(orig.Leaf);
            nodeMap[orig] = copy;
            if (orig.Children != null)
            {
                foreach (var v in orig.Children)
                {
                    copy.AddChild(CopyTree(v, nodeMap));
                }
            }
            return copy;
        }


        private Module AddModule(Tree t)
        {
            var m = new Module();
            m.TreeNode = t;
            treeModuleMap[m.TreeNode] = m;
            return m;
        }

        /// <summary>
        /// Get the module containing the ith node, by inefficient traversal 
        /// of tree leaves.
        /// </summary>
        /// <param name="i"></param>
        /// <returns>module containing the ith node</returns>
        public Module GetModule(int i)
        {
            foreach (var tm in treeModuleMap)
            {
                if (tm.Key.Leaf == i)
                {
                    // two cases: either it's a singleton module if the node
                    // has never been merged, or it's the module of the parent
                    // tree node
                    if (tm.Key.Parent == tree)
                    {
                        return tm.Value;
                    }
                    else
                    {
                        return treeModuleMap[tm.Key.Parent];
                    }
                }
            }
            return null;
        }

        /// <summary>
        /// Create flat module hierarchy for graph over n
        /// nodes defined by the given edge pairs.
        /// </summary>
        /// <param name="n"></param>
        /// <param name="edges"></param>
        public Modules(int n, Tuple<int, int>[] edges)
        {
            leafCount = n;
            tree = new Tree();
            var ms = new Module[n];
            for (int i = 0; i < n; ++i)
            {
                var t = new Tree(i);
                ms[i] = AddModule(t);
                tree.AddChild(t);
            }
            foreach (var e in edges)
            {
                var u = ms[e.Item1];
                var v = ms[e.Item2];
                u.Outgoing.Add(v);
                v.Incoming.Add(u);
            }
        }

        /// <summary>
        /// Merge constructor: construct clone of m with modules 
        /// a and b merged.
        /// </summary>
        /// <param name="m"></param>
        /// <param name="a"></param>
        /// <param name="b"></param>
        public Modules(Modules m, Module a, Module b) 
        {
            Debug.Assert(a.TreeNode.Parent == b.TreeNode.Parent);
            Debug.Assert(a.TreeNode.PathFromRoot().First() == m.tree);
            this.leafCount = m.leafCount;

            // mappings from old tree nodes to new tree nodes
            var treeMap = new Dictionary<Tree, Tree>();
            tree = CopyTree(m.tree, treeMap);

            // mappings from old modules to new modules
            var moduleMap = new Dictionary<Module, Module>();
            foreach (var v in m.treeModuleMap.Values)
            {
                var mm = AddModule(treeMap[v.TreeNode]);
                moduleMap[v] = mm;
            }
            foreach (var v in moduleMap)
            {
                v.Value.Incoming.UnionWith(v.Key.Incoming.Select(i => moduleMap[i]));
                v.Value.Outgoing.UnionWith(v.Key.Outgoing.Select(o => moduleMap[o]));
            }

            Module ma = treeModuleMap[treeMap[a.TreeNode]];
            Module mb = treeModuleMap[treeMap[b.TreeNode]];
            edgesRemoved = Merge(ma, mb);
        }
        public int edgesRemoved = 0;

        //private void MergeLeaf(Module a, Module b, Tree parent)
        //{
        //    Debug.Assert(!a.IsLeaf && b.IsLeaf);
        //    MergeNeighbours(a, b, m => m.Outgoing, m => m.Incoming);
        //    MergeNeighbours(a, b, m => m.Incoming, m => m.Outgoing);
        //    parent.RemoveChild(b.TreeNode);
        //    a.TreeNode.AddChild(b.TreeNode);
        //}

        //private void MergeNeighbours(Module a, Module b, Func<Module, ISet<Module>> getIn, Func<Module, ISet<Module>> getOut)
        //{
        //    // merge leaf module b into group module a
        //    Debug.Assert(!a.IsLeaf && b.IsLeaf);
        //    var aOutgoing = getOut(a);
        //    var bOutgoing = getOut(b);
        //    var aIncoming = getIn(a);
        //    var bIncoming = getIn(b);
        //    var outgoingCommon = aOutgoing.Intersect(bOutgoing).ToArray();
        //    var children = a.TreeNode.Children.Select(c => treeModuleMap[c]);

        //    // if edge from b to a then replace with edge from b to each member of a
        //    if (bOutgoing.Contains(a))
        //    {
        //        bOutgoing.Remove(a);
        //        aIncoming.Remove(b);
        //        bOutgoing.UnionWith(children);
        //        foreach (var v in children)
        //        {
        //            getIn(v).Add(b);
        //        }
        //    }

        //    // if a.Outgoing is not a subset of b.Outgoing then we need to reinstate all a.outgoing - outgoingCommon on a's children
        //    var mismatch = aOutgoing.Except(outgoingCommon.Union(new[] {b})).ToArray();
        //    if (mismatch.Any())
        //    {
        //        aOutgoing.ExceptWith(mismatch);
        //        foreach (var v in mismatch)
        //        {
        //            getIn(v).Remove(a);
        //            getIn(v).UnionWith(children);
        //        }
        //        foreach (var v in children)
        //        {
        //            getOut(v).UnionWith(mismatch);
        //        }
        //    }

        //    // b.Outgoing that are also in outgoingCommon are now redundant so remove
        //    bOutgoing.ExceptWith(outgoingCommon);
        //    foreach (var v in outgoingCommon)
        //    {
        //        getIn(v).Remove(b);
        //    }
        //}

        /// <summary>
        /// When a new module mm is created as a parent of modules a and b,
        /// any common edges of a and b (in a given direction) can be promoted
        /// to mm.
        /// </summary>
        /// <param name="mm"></param>
        /// <param name="a"></param>
        /// <param name="b"></param>
        /// <param name="getOut"></param>
        /// <param name="getIn"></param>
        private int PromoteEdges(Module mm, Module a, Module b, Func<Module, ISet<Module>> getOut, Func<Module, ISet<Module>> getIn)
        {
            var aOut = getOut(a);
            var bOut = getOut(b);
            var mmOut = getOut(mm);
            var outCommon = aOut.Intersect(bOut);
            int edgesRemoved = 0;

            // add outCommon to mmOut and rewire other ends of edges
            foreach (var v in outCommon)
            {
                var vIn = getIn(v);
                vIn.Remove(a);
                vIn.Remove(b);
                vIn.Add(mm);
                mmOut.Add(v);
                ++edgesRemoved;
            }

            aOut.ExceptWith(mmOut);
            bOut.ExceptWith(mmOut);
            return edgesRemoved;
        }

        public int TestMerge(Module a, Module b)
        {
            int removeEdges = a.Outgoing.Intersect(b.Outgoing).Count();
            removeEdges += a.Outgoing.Intersect(b.Outgoing).Count();
            return removeEdges;
        }

        /// <summary>
        /// Merge modules a and b into a new module under their common parent.
        /// Then flatten if necessary.  That is, if---as a result of promotion 
        /// of edges to the new module---a has no 
        /// more edges, then a is removed.  Similar for b.
        /// </summary>
        /// <param name="a">module to merge</param>
        /// <param name="b">module to merge</param>
        public int Merge(Module a, Module b)
        {
            this.cost = -1;
            this.edges = -1;
            Tree parent = a.TreeNode.Parent;
            Debug.Assert(parent == b.TreeNode.Parent);

            parent.RemoveChild(a.TreeNode);
            parent.RemoveChild(b.TreeNode);
            var mm = AddModule(new Tree(a.TreeNode, b.TreeNode));
            parent.AddChild(mm.TreeNode);
            int edgesRemoved = 0;
            edgesRemoved += PromoteEdges(mm, a, b, m => m.Outgoing, m => m.Incoming);
            edgesRemoved += PromoteEdges(mm, a, b, m => m.Incoming, m => m.Outgoing);

            // if a.incoming and a.outgoing are empty, remove a and promote its children
            // similar for b
            Flatten(a, mm);
            Flatten(b, mm);

            return edgesRemoved;
        }

        private void Flatten(Module module, Module parent)
        {
            if (module.TreeNode.Children != null && module.Outgoing.Count == 0 && module.Incoming.Count == 0)
            {
                parent.TreeNode.RemoveChild(module.TreeNode);
                parent.TreeNode.AddChildren(module.TreeNode.Children);
                treeModuleMap.Remove(module.TreeNode);
            }
        }

        IEnumerator<Tuple<Tree, Tree>> mergeTraversalPosition;
        public void InitializeMergeTraversal()
        {
            mergeTraversalPosition = new TopLevelMergeTraversalPosition(tree);
        }

        public Tuple<Module, Module> GetNextMergePair()
        {
            if (mergeTraversalPosition.MoveNext())
            {
                var m = mergeTraversalPosition.Current;
                return Tuple.Create(treeModuleMap[m.Item1], treeModuleMap[m.Item2]);
            }
            return null;
        }

        public IEnumerable<Tuple<int, int>> GetFlatGraph()
        {
            var edges = new List<Tuple<int, int>>();
            foreach (var l in tree.LeafNodes)
            {
                int u = l.Leaf.Value;
                foreach (var t in l.PathToRoot().Where(t => t != tree))
                {
                    foreach (var m in treeModuleMap[t].Outgoing)
                    {
                        foreach (var v in m.TreeNode.Leaves)
                        {
                            edges.Add(Tuple.Create(u, v.Value));
                        }
                    }
                }
            }
            return edges;
        }

        private class MergeTraversalPosition : IEnumerator<Tuple<Tree, Tree>>
        {
            LinkedListNode<Tree> firstMergeCandidate;
            LinkedListNode<Tree> secondMergeCandidate;
            IEnumerator<Tree> treeEnumerator;
            Tree root;

            public MergeTraversalPosition(Tree root)
            {
                this.root = root; 
                treeEnumerator = root.GetEnumerator();
            }

            public bool MoveNext()
            {
                if (firstMergeCandidate == null)
                {
                    while (treeEnumerator.MoveNext())
                    {
                        var node = treeEnumerator.Current;
                        if (node.Children != null && node.Children.Count > 1)
                        {
                            firstMergeCandidate = node.Children.First;
                            secondMergeCandidate = node.Children.First.Next;
                            return true;
                        }
                    }
                    return false;
                }
                if (secondMergeCandidate.Next != null)
                {
                    secondMergeCandidate = secondMergeCandidate.Next;
                    return true;
                }
                if (firstMergeCandidate.Next.Next != null)
                {
                    firstMergeCandidate = firstMergeCandidate.Next;
                    secondMergeCandidate = firstMergeCandidate.Next;
                    return true;
                }
                firstMergeCandidate = null;
                return MoveNext();
            }

            public Tuple<Tree, Tree> Current
            {
                get
                {
                    if (firstMergeCandidate == null)
                    {
                        throw new InvalidOperationException("Attempt to call current before MoveNext or when no more items in enumerator");
                    }
                    return Tuple.Create(firstMergeCandidate.Value, secondMergeCandidate.Value);
                }
            }

            public void Dispose()
            {
            }

            object IEnumerator.Current
            {
                get { return this.Current; }
            }

            public void Reset()
            {
                treeEnumerator = root.GetEnumerator();
                firstMergeCandidate = null;
            }
        }

        private class TopLevelMergeTraversalPosition : IEnumerator<Tuple<Tree, Tree>>
        {
            LinkedListNode<Tree> firstMergeCandidate;
            LinkedListNode<Tree> secondMergeCandidate;
            LinkedList<Tree> children;
            Tree root;

            public TopLevelMergeTraversalPosition(Tree root)
            {
                this.root = root;
                children = root.Children;
            }

            public bool MoveNext()
            {
                if (firstMergeCandidate == null)
                {
                    if (children != null && children.Count > 1)
                    {
                        firstMergeCandidate = children.First;
                        secondMergeCandidate = children.First.Next;
                        return true;
                    }
                    return false;
                }
                if (secondMergeCandidate.Next != null)
                {
                    secondMergeCandidate = secondMergeCandidate.Next;
                    return true;
                }
                if (firstMergeCandidate.Next.Next != null)
                {
                    firstMergeCandidate = firstMergeCandidate.Next;
                    secondMergeCandidate = firstMergeCandidate.Next;
                    return true;
                }
                firstMergeCandidate = null;
                return false;
            }

            public Tuple<Tree, Tree> Current
            {
                get
                {
                    if (firstMergeCandidate == null)
                    {
                        throw new InvalidOperationException("Attempt to call current before MoveNext or when no more items in enumerator");
                    }
                    return Tuple.Create(firstMergeCandidate.Value, secondMergeCandidate.Value);
                }
            }

            public void Dispose()
            {
            }

            object IEnumerator.Current
            {
                get { return this.Current; }
            }

            public void Reset()
            {
                firstMergeCandidate = null;
            }
        }

        public override string ToString()
        {
            return ToString(tree);
        }

        /// <summary>
        /// ordered traversal of the tree to obtain unique id
        /// </summary>
        /// <param name="t"></param>
        /// <returns></returns>
        private string ToString(Tree t)
        {
            string s = "";
            if (t.Children != null)
            {
                var children = new List<string>(t.Children.Select(c=>ToString(c)));
                children.Sort();
                s += "(" + children[0];
                for (int i = 1; i < children.Count; ++i)
                {
                    s += "," + children[i];
                }
                s += ")";
            }
            if (t.Leaf != null)
            {
                s += t.Leaf.ToString();
            }
            return s;
        }
    }

    public class KSmallestQueue<TKey, TValue> : IEnumerable<TValue>
        where TKey : IComparable
    {
        readonly MaxHeap<KeyValuePair<TKey, TValue>> heap;
        readonly int k;

        public KSmallestQueue(int k)
        {
            this.k = k;
            var comparer = Comparer<KeyValuePair<TKey, TValue>>.Create((x, y) => x.Key.CompareTo(y.Key));
            this.heap = new MaxHeap<KeyValuePair<TKey, TValue>>(comparer);
        }

        public bool InsertIfBetter(TKey key, Func<TValue> getValue)
        {
            bool inserted = false;
            if (heap.Count < k || key.CompareTo(heap.Top().Key) < 0)
            {
                heap.Add(new KeyValuePair<TKey, TValue>(key, getValue()));
                inserted = true;
            }
            if (heap.Count > k)
            {
                heap.ExtractDominating();
            }
            return inserted;
        }

        public int K
        {
            get { return k; }
        }

        public int Count
        {
            get { return heap.Count; }
        }

        public KeyValuePair<TKey, TValue> Pop()
        {
            return heap.ExtractDominating();
        }

        public KeyValuePair<TKey, TValue> Worst()
        {
            return heap.Top();
        }

        public IEnumerator<TValue> GetEnumerator()
        {
            return heap.Select(t => t.Value).GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }
    }

    public class KSmallestUniqueQueue : IEnumerable<Modules>
    {
        readonly HashSet<string> lookup = new HashSet<string>();
        readonly MaxHeap<Modules> heap;
        int k;

        private void Insert(Modules m, string id) 
        {
            lookup.Add(id);
            heap.Add(m);
        }

        public KSmallestUniqueQueue(Modules m0, int k) : this(k)
        {
            Insert(m0, m0.ToString());
        }

        public KSmallestUniqueQueue(int k)
        {
            var comparer = Comparer<Modules>.Create((x, y) => x.Cost.CompareTo(y.Cost));
            this.heap = new MaxHeap<Modules>(comparer);
            this.k = k;
        }

        public int Count
        {
            get { return heap.Count; }
        }

        public int K
        {
            get { return k; }
        }

        /// <summary>
        /// Insert m into heap only if:
        ///  - a similar solution does not already exist in the heap
        ///  - and (heap size is less than k
        ///  - or m.Cost &lt the minimum of costBound and biggest cost currently in heap 
        /// </summary>
        /// <param name="m">a configuration of modules</param>
        /// <param name="costBound">a cost that m must beat to be inserted in heap</param>
        /// <returns>true if m was inserted</returns>
        public bool InsertIfDifferentAndBetter(Modules m, int costBound = int.MaxValue)
        {
            var mID = m.ToString();
            if (lookup.Contains(mID))
            {
                return false;
            }
            if (heap.Count < k || m.Cost < Math.Min(costBound, heap.Top().Cost))
            {
                Insert(m, mID);
                if (heap.Count > k)
                {
                    heap.ExtractDominating();
                }
                return true;
            }
            return false;
        }

        public IEnumerator<Modules> GetEnumerator()
        {
            return heap.GetEnumerator();
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }
    }

    public class BeamSearch
    {
        KSmallestUniqueQueue beam;
        public BeamSearch(Modules m0, int k)
        {
            beam = new KSmallestUniqueQueue(m0, k);
        }

        public int BeamCount
        {
            get { return beam.Count; }
        }

        public IEnumerable<Modules> Beam
        {
            get
            {
                return beam;
            }
        }

        public Modules Best
        {
            get
            {
                var ms = from m in beam
                         orderby m.Cost
                         select m;
                return ms.First();
            }
        }

        public bool Improve()
        {
            bool improved = false;
            if (Modules.CrossingWeight == 0)
            {
                Tuple<Module, Module> p;
                var best = new List<Tuple<int, Modules, Tuple<Module, Module>>>();
                foreach (var m in beam.ToArray())
                {
                    m.InitializeMergeTraversal();
                    while ((p = m.GetNextMergePair()) != null)
                    {
                        int edgesRemoved = m.TestMerge(p.Item1, p.Item2);
                        best.Add(Tuple.Create(edgesRemoved, m, p));
                    }
                }
                var bestPairs = from b in best
                            orderby b.Item1 descending
                            select b;
                foreach (var b in bestPairs.Take(beam.K))
                {
                    var m = b.Item2;
                    var mCost = m.Cost;
                    p = b.Item3;
                    var mprime = new Modules(m, p.Item1, p.Item2);
                    improved = beam.InsertIfDifferentAndBetter(mprime, mCost) || improved;
                }
            }
            else
            {
                foreach (var m in beam.ToArray())
                {
                    var mCost = m.Cost;
                    m.InitializeMergeTraversal();
                    Tuple<Module, Module> p;
                    while ((p = m.GetNextMergePair()) != null)
                    {
                        var mprime = new Modules(m, p.Item1, p.Item2);
                        Debug.Assert(m.EdgeCount == mprime.EdgeCount + mprime.edgesRemoved);

                        improved = beam.InsertIfDifferentAndBetter(mprime, mCost) || improved;
                    }
                }
            }
            return improved;
        }
    }

    public class GreedyMerge
    {
        Modules ms;
        public GreedyMerge(Modules ms)
        {
            this.ms = ms;
        }
        public bool Improve()
        {
            var edges0 = ms.EdgeCount;
            var bestImprovement = 0;
            Module bestA = null;
            Module bestB = null;
            ms.InitializeMergeTraversal();
            Tuple<Module, Module> p;
            while ((p = ms.GetNextMergePair()) != null)
            {
                int edgesRemoved = ms.TestMerge(p.Item1, p.Item2);
                if (edgesRemoved > bestImprovement)
                {
                    bestA = p.Item1;
                    bestB = p.Item2;
                    bestImprovement = edgesRemoved;
                }
            }
            if (bestImprovement > 0)
            {
                int edgesRemoved = ms.Merge(bestA, bestB);
                Debug.Assert(edgesRemoved == bestImprovement);
                Debug.Assert(edges0 == ms.EdgeCount + edgesRemoved);
                return true;
            }
            return false;
        }
    }
}
